import heapq
import time
import networkx as nx
import matplotlib.pyplot as plt
from generate import GenerateNetwork
import random

# Add:  监控网络中链路残余的密钥
#       self.G[u][v]["remain_key_num"] = 0 网络的链路中增加了剩余的密钥数量
#       监控某个SD对还没满足的密钥数量
#       self.sd_list 末尾处增加了已满足的密钥量


class QKDNetwork:
    """
    self.G[u][v]["remain_key_num"] = 0 网络的链路中增加了剩余的密钥数量
    self.sd_list 末尾处增加了已满足的密钥量
    """

    def sd_demand(self):
        tmp_rate = [1]
        if self.hete:
            tmp_rate = [2, 4]
        return random.choice(tmp_rate)  # QKD请求的速率随机从 2/4/6/8 Kbps 取

    def ps_rate(self):
        tmp_rate = [1]
        if self.hete:
            tmp_rate = [1, 2]
        return random.choice(tmp_rate)  # 密钥发射器速率随机从 1/2 Kbps 取

    def __init__(
        self,
        showTopology=False,
        num_node=20,
        sd_num=20,
        alpha=0.35,
        no_add_ps=False,
        hete=False,
        hybrid=False,
    ):
        self.showTopology = showTopology
        self.hete = hete
        self.no_add_ps = no_add_ps
        self.num_nodes = num_node  # 节点数量
        self.num_sd_pairs = sd_num  # 请求个数
        self.alpha = alpha  # Waxmax Alpha
        self.hybrid = hybrid  # 是否使用混合 relay 网络

        self.spd = []  # 光子源探测器的位置，(u,v)表示在链路uv的v端点有光子源
        self.mdi_links = []

        self.prepareNetwork()

        if self.hybrid:
            self.prepareMDILinks()

        # print("网络准备完成")
        # 展示网络
        if self.showTopology:
            self.showNetwork()

        # QKD请求的SD对
        self.sd_list = self.generate_random_sd_pairs(
            self.num_sd_pairs, 0, self.num_nodes - 1
        )

        # print("请求准备完成")
        if not no_add_ps:
            self.getCandidatePaths()
        # print("候选路径准备完成")
        if not no_add_ps:
            self.estimateMinTimeSlot()
        # print(f"预估时隙完成: {self.estimate_max_time_slot}")

    def prepareMDILinks(self):
        # 预设一系列MDI的链路，主要是 发射器 -> 接收器 <- 发射器
        preset_mdi_link_pair = []  # (m v n)
        for i in range(self.num_nodes):
            if self.G.nodes[i]["transmitter"] == 0:
                # 配备有接收器，从邻居节点中，找一对发射节点
                j_is_transmitter = []
                for j in self.G.neighbors(i):
                    if self.G.nodes[j]["transmitter"] == 1:
                        j_is_transmitter.append(j)
                if len(j_is_transmitter) >= 2:
                    # 从 j_is_transmitter 中随机挑两个点出来
                    arr = random.sample(j_is_transmitter, 2)
                    preset_mdi_link_pair.append((arr[0], i, arr[1]))
        # MDI-QKD链路的合集
        mdi_link_pair_count = len(preset_mdi_link_pair) // 3  # 取1/3
        if mdi_link_pair_count > 0:
            self.mdi_links = random.sample(preset_mdi_link_pair, mdi_link_pair_count)
        else:
            self.mdi_links = []
        # print("mdi_link_pair_count", mdi_link_pair_count)
        # print("mdi_links", self.mdi_links)
        # 将 MDI-Pair 节点以权重为2加入到新的网络中，并剔除相应的边
        tmpG = nx.Graph()
        for edge in self.G.edges:
            tmpG.add_edge(edge[0], edge[1])
        for mdi_link in self.mdi_links:
            tmpG.remove_edge(mdi_link[0], mdi_link[1])
            tmpG.remove_edge(mdi_link[1], mdi_link[2])
            tmpG.add_edge(mdi_link[0], mdi_link[2], weight=2)
        self.hybridNetworkX = tmpG

    def removePS(self):
        # 去掉所有的PS
        for node in self.G.nodes:
            self.G.nodes[node]["transmitter"] = 0

    def clearLoad(self):
        for u, v, _ in self.G.edges(data=True):
            self.G[u][v]["load"] = 0

    def getShortestPathWithTransmitter(
        self, source, target, forbidden_links: list = []
    ):
        # 获取网络中源目节点对的中继路径，要求穿过Transmitter
        processedNodes = []  # 存放已处理过的节点ID
        toBeProcessedNodes = []  # 待处理的节点ID

        for i in self.G.nodes:  # 初始化
            node = self.G.nodes[i]
            node["distance"] = 100000000000
            node["physicsDistance"] = 100000000000
            node["shortestPathLastNodeID"] = -1  # 路径的上一个节点ID为-1
            node["isAddedToQueue"] = False
            node["isVisited"] = False

        source_node = self.G.nodes[source]
        source_node["distance"] = (
            0 if source_node["transmitter"] == 1 else 1
        )  # 初始化距离，如果是transmitter就初始化为1，否则为0
        source_node["physicsDistance"] = 0
        source_node["shortestPathLastNodeID"] = (
            source  # 最短路径上自己的上一个节点就是自己
        )
        toBeProcessedNodes.append(source)
        source_node["isAddedToQueue"] = True

        while len(toBeProcessedNodes) != 0:
            toBeProcessedNodes = sorted(
                toBeProcessedNodes, key=lambda x: self.G.nodes[x]["physicsDistance"]
            )  # 排序
            tmp = toBeProcessedNodes.pop(0)
            processedNodes.append(tmp)
            self.G.nodes[tmp]["isVisited"] = True
            for i in self.G.neighbors(tmp):
                if ([tmp, i] in forbidden_links) or ([i, tmp] in forbidden_links):
                    continue  # 去掉断掉的边
                node_adjacent = self.G.nodes[i]
                node_tmp = self.G.nodes[tmp]
                if node_adjacent["isVisited"] or (
                    node_adjacent["transmitter"] == 0 and node_tmp["distance"] + 1 > 1
                ):
                    continue
                if (
                    node_adjacent["distance"] == 100000000000
                    or node_adjacent["physicsDistance"]
                    > node_tmp["physicsDistance"] + 1
                ):
                    node_adjacent["distance"] = node_tmp["distance"] + 1
                    node_adjacent["physicsDistance"] = node_tmp["physicsDistance"] + 1
                    node_adjacent["shortestPathLastNodeID"] = tmp
                if node_adjacent["transmitter"] == 1:
                    node_adjacent["distance"] = 0
                if not node_adjacent["isAddedToQueue"]:
                    toBeProcessedNodes.append(i)
                    node_adjacent["isAddedToQueue"] = True
        try:
            path = [target]
            while target != source and target != None:
                tmpLastID = self.G.nodes[target]["shortestPathLastNodeID"]
                path.insert(0, tmpLastID)
                target = tmpLastID
        except Exception as _:
            return []  # 说明没有合适的路

        # print(path)
        # print([self.G.nodes[j]["transmitter"] for j in path])
        return path

    def getEntirePathWithMDI(self, fold_path: list = []):
        if len(fold_path) == 0:
            return []
        ret = [fold_path[0]]
        for i in range(len(fold_path) - 1):
            start, end = fold_path[i], fold_path[i + 1]
            flag = False
            for mdi_link in self.mdi_links:
                if mdi_link[0] == start and mdi_link[2] == end:
                    ret.append(mdi_link[1])
                    ret.append(mdi_link[2])
                    flag = True
                    break
                elif mdi_link[2] == start and mdi_link[0] == end:
                    ret.append(mdi_link[1])
                    ret.append(mdi_link[0])
                    flag = True
                    break
            if not flag:
                ret.append(end)
        if len(fold_path) != len(ret):
            print(fold_path, " -> ", ret)
        return ret

    def getSPWithMDI(self, source, target, forbidden_links: list = []):
        sp = nx.shortest_path(self.hybridNetworkX, source, target, weight="weight")
        return self.getEntirePathWithMDI(sp)

    def getSPWithTransmitterWithNetworkX(
        self, source, target, forbidden_links: list = []
    ):
        tmpG = nx.Graph()
        for edge in self.G.edges:
            if edge in forbidden_links or [edge[1], edge[0]] in forbidden_links:
                continue
            if (
                self.G.nodes[edge[0]]["transmitter"] == 1
                or self.G.nodes[edge[1]]["transmitter"] == 1
            ):
                tmpG.add_edge(edge[0], edge[1])
        return nx.shortest_path(tmpG, source, target, weight="weight")

    def getCandidatePathsWithSPD(self):
        if len(self.spd) == 0:
            return
        candidate_paths = []
        self.max_candidate_size = 0
        self.max_candidate_path_length = 0
        # 构建无向图
        tmpG = nx.Graph()
        for i in self.spd:
            tmpG.add_edge(i[0], i[1])

        for sd in self.sd_list:
            # 计算候选路径
            paths = []
            K = 5
            for i, path in enumerate(nx.shortest_simple_paths(tmpG, sd[0], sd[1])):
                self.max_candidate_path_length = (
                    self.max_candidate_path_length
                    if self.max_candidate_path_length > len(path)
                    else len(path)
                )
                paths.append(path)
                if i == K:  # 当找到 K 条路径时停止
                    break
            self.max_candidate_size = (
                self.max_candidate_size
                if self.max_candidate_size > len(paths)
                else len(paths)
            )
            candidate_paths.append({"sd": sd, "paths": paths})
        self.candidate_paths_spd = candidate_paths
        return

    def getCandidatePathsWithMDI(self):
        if len(self.mdi_links) == 0:
            return
        candidate_paths = []
        self.max_candidate_size = 0
        self.max_candidate_path_length = 0
        # 构建无向图
        tmpG = self.hybridNetworkX
        for sd in self.sd_list:
            # 计算候选路径
            paths = []
            K = 5
            for i, path in enumerate(
                nx.shortest_simple_paths(tmpG, sd[0], sd[1], weight="weight")
            ):
                entire_path = self.getEntirePathWithMDI(path)
                len_path = len(entire_path)
                self.max_candidate_path_length = max(
                    self.max_candidate_path_length, len_path
                )
                paths.append(entire_path)
                if i == K:  # 当找到 K 条路径时停止
                    break
            self.max_candidate_size = max(self.max_candidate_size, len(paths))
            candidate_paths.append({"sd": sd, "paths": paths})
        self.candidate_paths_mdi = candidate_paths

    def getCandidatePaths(self):
        candidate_paths = []
        self.max_candidate_size = 0
        self.max_candidate_path_length = 0
        for sd in self.sd_list:
            # 计算候选路径
            paths = []
            path_1 = self.getShortestPathWithTransmitter(sd[0], sd[1])
            self.max_candidate_path_length = (
                self.max_candidate_path_length
                if self.max_candidate_path_length > len(path_1)
                else len(path_1)
            )
            paths.append(path_1)
            for i in range(len(path_1) - 1):
                link = [path_1[i], path_1[i + 1]]
                # 逐一断开
                yen_path = self.getShortestPathWithTransmitter(
                    sd[0], sd[1], forbidden_links=[link]
                )
                if len(yen_path) != 0 and yen_path not in paths:
                    paths.append(yen_path)
                    self.max_candidate_path_length = (
                        self.max_candidate_path_length
                        if self.max_candidate_path_length > len(yen_path)
                        else len(yen_path)
                    )
            if len(path_1) > 2:
                for i in range(len(path_1) - 2):
                    link_1 = [path_1[i], path_1[i + 1]]
                    link_2 = [path_1[i + 1], path_1[i + 2]]
                    # 逐二断开
                    yen_path = self.getShortestPathWithTransmitter(
                        sd[0], sd[1], forbidden_links=[link_1, link_2]
                    )
                    if len(yen_path) != 0 and yen_path not in paths:
                        paths.append(yen_path)
                        self.max_candidate_path_length = (
                            self.max_candidate_path_length
                            if self.max_candidate_path_length > len(yen_path)
                            else len(yen_path)
                        )
            # print("Yen", paths)
            self.max_candidate_size = (
                self.max_candidate_size
                if self.max_candidate_size > len(paths)
                else len(paths)
            )
            candidate_paths.append({"sd": sd, "paths": paths})
        self.candidate_paths = candidate_paths
        return

    def estimateMinTimeSlot(self):
        # 估计最小的完成时间隙数量，每个请求都按照最短路
        def addFlowOnPath(sd):  # 施加流
            path = self.getShortestPathWithTransmitter(sd[0], sd[1])
            for i in range(len(path) - 1):
                link_start = path[i]
                link_end = path[i + 1]
                self.G[link_start][link_end]["load"] += sd[2]

        for tmp_sd in self.sd_list:
            addFlowOnPath(tmp_sd)

        max_time_slot = 0
        while True:
            flag = True
            for i in self.G.nodes:  # 初始化
                node = self.G.nodes[i]
                if node["transmitter"] == 1:
                    # 选择一个需要密钥最多的链路进行服务
                    max_rate_link_node = 0
                    max_rate = 0
                    for neighbor in self.G.neighbors(i):
                        if self.G[i][neighbor]["load"] > max_rate:
                            max_rate = self.G[i][neighbor]["load"]
                            max_rate_link_node = neighbor
                            flag = False
                    if max_rate > 0:
                        self.G[max_rate_link_node][i]["load"] -= node[
                            "transmitter_rate"
                        ]
            if flag:
                break
            else:
                max_time_slot += 1
        self.estimate_max_time_slot = max_time_slot

    def generate_random_sd_pairs(self, num_pairs, min_val, max_val):
        pairs = set()
        arr_mdi_relay_id = [i[1] for i in self.mdi_links]
        while len(pairs) < num_pairs:
            sd_start_node = random.choice(
                [x for x in range(min_val, max_val + 1) if x not in arr_mdi_relay_id]
            )
            sd_end_node = random.choice(
                [x for x in range(min_val, max_val + 1) if x not in arr_mdi_relay_id]
            )
            sd_request_demand = self.sd_demand()
            if sd_start_node != sd_end_node:
                # 在末尾处增加了一个已满足的密钥量
                pairs.add((sd_start_node, sd_end_node, sd_request_demand, 0))
        return list(pairs)

    def showNetwork(self):
        pos = nx.spring_layout(self.G, seed=42)
        node_colors = [
            "red" if self.G.nodes[node]["transmitter"] == 1 else "skyblue"
            for node in self.G.nodes
        ]
        plt.figure(figsize=(8, 6))
        nx.draw(
            self.G,
            pos,
            with_labels=True,
            node_color=node_colors,
            edge_color="gray",
            node_size=500,
            font_size=10,
        )
        plt.title("Random Graph with Transmitters")
        plt.show()

    def prepareNetwork(self):
        beta = 0.1
        # self.G: nx.Graph = nx.waxman_graph(self.num_nodes, self.alpha, beta)
        self.G: nx.Graph = GenerateNetwork().generate_network(
            self.num_nodes, self.alpha, beta
        )
        # 检查图是否连通
        if not nx.is_connected(self.G):
            # 获取所有连通分量
            components = list(nx.connected_components(self.G))
            # 连接每个连通分量
            for i in range(len(components) - 1):
                # 在第i个和第i+1个连通分量之间添加一条边
                # 选择每个连通分量中的任意一个节点
                a = next(iter(components[i]))
                b = next(iter(components[i + 1]))
                self.G.add_edge(a, b)

        # 准备网络，使得网络链路的两个端点中至少有一个为发射器节点
        for node in self.G.nodes():
            self.G.nodes[node]["transmitter"] = 0
            self.G.nodes[node]["transmitter_rate"] = 0

        for u, v, _ in self.G.edges(data=True):
            self.G[u][v]["load"] = 0
            # 增加一个残余的密钥量
            self.G[u][v]["remain_key_num"] = 0

        if self.no_add_ps:
            return

        def hasFeasiblePath(u, v) -> bool:
            max_length = 2

            def isFeasiblePath(path: list) -> bool:
                for tmpLink in [[path[i], path[i + 1]] for i in range(len(path) - 1)]:
                    link_start_transmitter = self.G.nodes[tmpLink[0]]["transmitter"]
                    link_end_transmitter = self.G.nodes[tmpLink[1]]["transmitter"]
                    if link_start_transmitter == 0 and link_end_transmitter == 0:
                        return False
                return True

            for path in nx.all_simple_paths(
                self.G, source=u, target=v, cutoff=max_length
            ):
                if len(path) <= max_length + 1:
                    if isFeasiblePath(path):
                        return True
            return False

        for _ in range(100):
            for u in self.G.nodes:
                if (
                    self.G.nodes[u]["transmitter"] == 1
                    and self.G.nodes[v]["transmitter"] == 1
                ):
                    # 两边都配备有密钥发生器，随机去掉一个
                    random_node = random.choice([u, v])
                    self.G.nodes[random_node]["transmitter"] = 0
            for u, v in self.G.edges:
                if (
                    self.G.nodes[u]["transmitter"] == 0
                    and self.G.nodes[v]["transmitter"] == 0
                ):
                    # if((not hasFeasiblePath(u, v))):
                    # 两边都没有配备有密钥发生器且无中转节点，添加一个度数较大的
                    smaller_degree_node = (
                        u if self.G.degree(u) < self.G.degree(v) else v
                    )
                    self.G.nodes[smaller_degree_node]["transmitter"] = 1

        trans_count = 0
        not_trans_count = 0
        for i in self.G.nodes:
            if self.G.nodes[i]["transmitter"] == 1:
                trans_count += 1
                self.G.nodes[i]["transmitter_rate"] = self.ps_rate()
            else:
                not_trans_count += 1

        # print(f"Transmitter: {trans_count}")
        # print(f"NotTransmitter: {not_trans_count}")
        # print(f"发射器编排完成")
